#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import asyncio
import os
from datetime import datetime
import json
import sys
from typing import List, Optional

from src.agents.AbstractGeneratorAgent import AbstractGeneratorAgent
from src.agents.KeywordExtractorAgent import KeywordExtractorAgent


from src.agents.OutlineBuilderAgent import OutlineBuilderAgent
from src.agents.CaseAnalyzerAgent import CaseAnalyzerAgent
from src.agents.PolicyEvaluatorAgent import PolicyEvaluatorAgent
from src.agents.RetrieverAgent import RetrieverAgent, Article
from src.agents.base_config import StateManager
from src.utils.env_config import get_task_name, file_resume, file_report

def deal_articles(articles: List[Article]) -> str:
    summary = ""
    for article in articles:
        summary += f"{article["title"]}\n{article["content"]}\n"
    return summary

class TaskOrchestrator3:
    """增强型主控智能体：支持动态任务调整与上下文保持"""

    def __init__(self):
        self.state_manager = StateManager()
        self.keyword_agent = KeywordExtractorAgent()
        self.literature_agent = RetrieverAgent()
        self.abstract_agent = AbstractGeneratorAgent()
        self.outline_agent = OutlineBuilderAgent()
        self.case_agent = CaseAnalyzerAgent()
        self.policy_agent = PolicyEvaluatorAgent()
        self.last_task_id = 0

    @staticmethod
    def should_case_analysis(retrieval_summary: str) -> bool:
        if retrieval_summary.find("技术突破") >= 0 and retrieval_summary.find("创新案例") >= 0:
            return True
        return False
    @staticmethod
    def should_policy_evaluation(retrieval_summary: str) -> bool:
        if retrieval_summary.find("政策法规") >= 0:
            return True
        return False

    @staticmethod
    def deal_resume():
        resume = False
        if not os.path.exists(file_resume):
            with open(file_resume, 'w', encoding='utf-8') as f:
                f.write('')  # 创建空文件
                f.close()
                return 0
        last_line =""
        with open(file_resume, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            if len(lines) > 0:
                last_line = lines[-1].rstrip('\n')  # 去除末尾换行符
                if last_line.find("[中断时间]") >= 0:
                    resume = True
            f.close()

        if resume:
            task = last_line[last_line.find("T"):last_line.find("T")+2]
            task_id = int(task[1:])
            current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            resume_log = f"[续跑时间] {current_time}\n"
            print(resume_log)
            with open(file_resume, 'a', encoding='utf-8') as f:
                f.write(resume_log)
                f.close()

            if task_id == 3:
                task_id = 7
            if task_id == 4:
                return 0

            return task_id
        else:
            return 0

    def write_resume(self, task):
        with open(file_resume, 'a', encoding='utf-8') as f:
            recover_log = f"[恢复执行] 从 {task}（{get_task_name(task)}）开始\n"
            print(recover_log)
            f.write(recover_log)
            f.close()


    async def execute_task(self, topic: str) -> None:
        """执行完整的报告生成任务"""
        print(f"开始执行报告生成任务：{topic}")
        """获取开始任务ID"""

        task_id = self.deal_resume()
        resume = False
        if task_id > 0:
            resume = True
        if task_id == 0:
            self.state_manager.clear_context()

        # 保存主题到上下文
        self.state_manager.save_context("topic", topic)
        start_resume = False
        if task_id <= 1:
            if resume:
                start_resume = True
                resume = False
            else :
                # T1: 关键词提取
                await self._execute_subtask("T1", self._task1_keyword_extraction, topic)
                self.last_task_id = 1

        if task_id <= 2:
            if start_resume:
                self.write_resume("T2")
                start_resume = False
            if resume:
                start_resume = True
                resume = False

            else:
                # T2: 文献检索
                keywords = self.state_manager.get_context("keywords")
                await self._execute_subtask("T2", self._task2_literature_retrieval, keywords)
                self.last_task_id = 2

        # 决策分析：确定后续任务流程
        retrieval_summary = deal_articles(self.state_manager.get_context("retrieval_summary"))
        if task_id <= 5:
            if self.should_case_analysis(retrieval_summary):
                # T5: 技术案例分析
                if start_resume:
                    self.write_resume("T5")
                    start_resume = False
                if resume:
                    start_resume = True
                    resume = False
                else:
                    await self._execute_subtask("T5", self._task5_case_analysis)
                    self.last_task_id = 5

        if task_id <= 6:
            if self.should_policy_evaluation(retrieval_summary):
                # T6: 政策影响评估
                if start_resume:
                    self.write_resume("T6")
                    start_resume = False
                if resume:
                    start_resume = True
                    resume = False
                await self._execute_subtask("T6", self._task6_policy_evaluation)
                self.last_task_id = 6

        if task_id <= 7:
            # T3: 摘要生成
            if start_resume:
                self.write_resume("T3")
                start_resume = False
            if resume:
                start_resume = True
                resume = False
            else:
                await self._execute_subtask("T3", self._task3_abstract_generation, retrieval_summary)
                self.last_task_id = 3

        if task_id <= 8:
            # T4: 报告框架搭建
            if start_resume:
                self.write_resume("T4")
                start_resume = False
            if resume:
                start_resume = True
                resume = False
            else:
                abstract = self.state_manager.get_context("abstract")
                await self._execute_subtask("T4", self._task4_outline_building, topic, abstract)
                self.last_task_id = 4

        print("=== 所有任务执行完成！ ===")

    async def _execute_subtask(self, task_id: str, task_func, *args) -> None:
        """执行子任务"""
        task_name = get_task_name(task_id)
        #清空子任务之前的状态
        self.state_manager.clear_task_status(task_id)
        self.state_manager.update_task_status(task_id, "未开始")
        print(f"\n开始执行 {task_id}: {task_name}")
        # 更新状态为"进行中"
        self.state_manager.update_task_status(task_id, "进行中")
        try:
            # 执行任务
            await task_func(*args)
            # 更新状态为"完成"
            self.state_manager.update_task_status(task_id, "完成")
            print(f"{task_id}: {task_name} 执行完成")

        except Exception as e:
            print(f"{task_id}: {task_name} 执行失败: {e}")
            self.state_manager.update_task_status(task_id, "失败")
            raise


    async def _task1_keyword_extraction(self, topic: str) -> None:
        """T1: 关键词提取"""
        keywords = await self.keyword_agent.extract_keywords(topic)
        self.state_manager.save_context("keywords", keywords)
        print(f"提取的关键词: {keywords}")

    async def _task2_literature_retrieval(self, keywords: List[str]) -> None:
        """T2: 文献检索"""
        retrieval_summary = await self.literature_agent.retrieve_literature(keywords)
        articles = self.literature_agent.articles
        self.state_manager.save_context("retrieval_summary", articles)
        print(f"文献检索摘要: {retrieval_summary[:100]}...")

    async def _task3_abstract_generation(self, retrieval_summary: str) -> None:
        """T3: 摘要生成"""
        case_analysis = self.state_manager.get_context("case_analysis")
        if case_analysis is not None and len(case_analysis) > 0:
            retrieval_summary += "\n技术案例分析:" + case_analysis
        policy_impact = self.state_manager.get_context("policy_impact")
        if policy_impact is not None and len(policy_impact) > 0:
            retrieval_summary += "\n政策影响评估:" + policy_impact

        abstract = await self.abstract_agent.generate_abstract(retrieval_summary)
        self.state_manager.save_context("abstract", abstract)
        print(f"生成的摘要: {abstract[:100]}...")

    async def _task4_outline_building(self, topic: str, abstract: str) -> None:
        """T4: 报告框架搭建"""
        markdown_content = await self.outline_agent.build_outline(topic, abstract)

        # 保存到文件
        with open(file_report, "w", encoding="utf-8") as f:
            f.write(markdown_content)

        self.state_manager.save_context("markdown_content", markdown_content)
        print("报告框架已生成并保存到 report_framework.md")

    async def _task5_case_analysis(self) -> None:
        """T5: 技术案例分析"""
        retrieval_summary = deal_articles(self.state_manager.get_context("retrieval_summary"))
        keywords = self.state_manager.get_context("keywords")
        abstract = self.state_manager.get_context("abstract")

        case_analysis = await self.case_agent.analyze_technical_cases(
            retrieval_summary, keywords, abstract
        )
        self.state_manager.save_context("case_analysis", case_analysis)
        print(f"技术案例分析: {case_analysis[:100]}...")

    async def _task6_policy_evaluation(self) -> None:
        """T6: 政策影响评估"""
        retrieval_summary = deal_articles(self.state_manager.get_context("retrieval_summary"))
        keywords = self.state_manager.get_context("keywords")
        abstract = self.state_manager.get_context("abstract")

        policy_impact = await self.policy_agent.evaluate_policy_impact(
            retrieval_summary, keywords, abstract
        )
        self.state_manager.save_context("policy_impact", policy_impact)
        print(f"政策影响评估: {policy_impact[:100]}...")

    def deal_exception(self):
        """处理异常,保存当然状态"""
        """
        追加写入resume_log.txt
        格式：
        [中断时间] 2025-10-15 14:20:00（最后完成任务：T2）
        [续跑时间] 2025-10-15 15:30:00
        [恢复执行] 从 T3（摘要生成）开始
        """
        task = f"T{self.last_task_id}"
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        resume_log = f"[中断时间] {current_time}（最后完成任务：{task}）\n"
        print(resume_log)
        with open(file_resume, 'a', encoding='utf-8') as f:
            f.write(resume_log)


def main():
    # """主函数"""
    # if len(sys.argv) != 2:
    #     print("使用方法: python task1_main.py '报告主题'")
    #     print("示例: python task1_agentscope.py '2025 Q3 AI 行业趋势'")
    #     return  '2025 Q3 AI大模型安全技术突破轻量化部署创新案例政策法规'
    # topic = sys.argv[1]
    topic = '2025 Q3 AI大模型安全技术突破轻量化部署创新案例政策法规'
    # 创建任务编排器并执行任务
    orchestrator = TaskOrchestrator3()
    try:
        asyncio.run(orchestrator.execute_task(topic))
    except KeyboardInterrupt:
        orchestrator.deal_exception()
        print(f"---任务已中断，可重新运行程序继续任务----")
        pass


if __name__ == "__main__":
    main()
